#!/usr/bin/python3

#Copyright (c) 2017 Sprint
################################################################################
# Licensed to the OpenAirInterface (OAI) Software Alliance under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The OpenAirInterface Software Alliance licenses this file to You under
# the terms found in the LICENSE file in the root of this source tree.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#-------------------------------------------------------------------------------
# For more information about the OpenAirInterface (OAI) Software Alliance:
#      contact@openairinterface.org
################################################################################
# file data_provisioning_users
# brief
# main author Sprint
# company Sprint
# email:
# other author Lionel GAUTHIER for python conversion

import argparse
import ipaddress
import logging
#import json
#import tempfile
from cassandra.cluster import Cluster

#---------------------------------------------------------------------
def str2bool(arg):
    ua = str(arg).lower()
    if 'true'.startswith(ua):
        return True
    elif 'false'.startswith(ua):
        return False
    else:
        return False

#---------------------------------------------------------------------
def ambr2json(br_ul, br_dl):
    ambr = '"AMBR":'

    extended_max_requested_bw_ul = 0
    extended_max_requested_bw_dl = 0

    if br_ul > 4294967295:
        # kilobits (1000 bits)
        extended_max_requested_bw_ul = int(br_ul/1000)
        br_ul = 4294967295

    ambr = ambr + '{'+'"Max-Requested-Bandwidth-UL":{0}'.format(br_ul)
      
    if br_dl > 4294967295:
        # kilobits (1000 bits)
        extended_max_requested_bw_dl = int(br_dl/1000)
        br_dl = 4294967295
    ambr = ambr + ',"Max-Requested-Bandwidth-DL":{0}'.format(br_dl)
        
    if extended_max_requested_bw_ul !=0:
        ambr = ambr + ',"Extended-Max-Requested-BW-UL":{0}'.format(extended_max_requested_bw_ul)
      
    if extended_max_requested_bw_dl != 0:
        ambr = ambr + ',"Extended-Max-Requested-BW-DL":{0}'.format(extended_max_requested_bw_dl)
        
    return ambr+'}'

#---------------------------------------------------------------------
def main():

    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)

    parser = argparse.ArgumentParser()
    parser.add_argument('-a', '--apn',          default='oai.ipv4',            help="default APN allowed for all IMSI of this HSS db")
    parser.add_argument('-A', '--apn2',         default='internet',            help="Non default APN allowed for all IMSI of this HSS db")
    parser.add_argument('-C', '--cassandra-cluster', default='127.0.0.1',      help="Cassandra list of nodes")
    
    parser.add_argument('-d', '--apn1-ambr-max-requested-bandwidth-dl', type=int, default=20000000, help="Max requested DL bandwidth in bits per seconds for APN1")
    parser.add_argument('-D', '--apn2-ambr-max-requested-bandwidth-dl', type=int, default=10000000, help="Max requested DL bandwidth in bits per seconds for APN2")
    parser.add_argument('-u', '--apn1-ambr-max-requested-bandwidth-ul', type=int, default=10000000, help="Max requested UL bandwidth in bits per seconds for APN1")
    parser.add_argument('-U', '--apn2-ambr-max-requested-bandwidth-ul', type=int, default=10000000, help="Max requested UL bandwidth in bits per seconds for APN2")
    parser.add_argument('-e', '--ue-ambr-max-requested-bandwidth-ul', type=int,   default=20000000, help="Max requested UL bandwidth in bits per seconds for UE")
    parser.add_argument('-E', '--ue-ambr-max-requested-bandwidth-dl', type=int,   default=40000000, help="Max requested DL bandwidth in bits per seconds for UE")
    
    parser.add_argument('-k', '--key',          default='8baf473f2f8fd09487cccbd7097c6862', help="LTE K (common for all IMSI)=8baf473f2f8fd09487cccbd7097c6862 for OP=11111111111111111111111111111111\n\
                                                                                                LTE K (common for all IMSI)=fec86ba6eb707ed08905757b1bb44b8f for OP=1006020f0a478bf6b699f15c062e42b3")
    parser.add_argument('-I', '--imsi-first',   default='001010000000001',     help="First IMSI for creating subscribers")
    parser.add_argument('-M', '--msisdn-first', default='+33663000021',        help="MSISDN of users (TODO: distincts MSISDN)")
    parser.add_argument('-m', '--mme-identity', default='mme.openair4G.eur',   help="Granted MME access to HSS")
    parser.add_argument('-n', '--no-of-users',  type=int, default=200,         help="Number of users")
    parser.add_argument('-o', '--opc',          default='8e27b6af0e692e750f32667a3b14605d', help="OPC (common for all IMSI)=8e27b6af0e692e750f32667a3b14605d for OP=11111111111111111111111111111111\n\
                                                                                                LTE K (common for all IMSI)=c42449363bbad02b66d16bc975d77cc1 for OP=1006020f0a478bf6b699f15c062e42b3")
    parser.add_argument('-q', '--sqn',          type=int, default=96,          help="SQN of the USIM of the subscriber")
    parser.add_argument('-r', '--realm',        default='openair4G.eur',       help="Realm of the CN")
    parser.add_argument('-s', '--static-ue-ipv4-allocation', default=None,     help="UE IPv4 address statically allocated by HSS, disabled if not provided")
    parser.add_argument('-S', '--static-ue-ipv6-allocation', default=None,     help="UE IPv6 address statically allocated by HSS, disabled if not provided")
    parser.add_argument('-t', '--truncate',     default='False',   choices=['True', 'TRUE', 'true', 'False', 'FALSE', 'false'],           help="Truncate tables, use with caution")
    parser.add_argument('-v', '--verbose',      default='True',    choices=['True', 'TRUE', 'true', 'False', 'FALSE', 'false'],          help="Display modified tables")
    args = parser.parse_args()


    # This will attempt to connection to a Cassandra instance on your local machine (127.0.0.1)
    # The set of IP addresses we pass to the :class:`~.Cluster` is simply an initial set of contact points.
    # After the driver connects to one of these nodes it will automatically discover the rest of the nodes
    # in the cluster and connect to them, so you don't need to list every node in your cluster.
    cluster = Cluster([args.cassandra_cluster])

    session = cluster.connect()
    # session.set_keyspace('mykeyspace')

    if str2bool(args.truncate):
        session.execute("""TRUNCATE vhss.msisdn_imsi ;""")
        session.execute("""TRUNCATE vhss.users_imsi ;""")

    apn1_ambr = ambr2json( args.apn1_ambr_max_requested_bandwidth_ul, args.apn1_ambr_max_requested_bandwidth_dl)
    apn2_ambr = ambr2json( args.apn2_ambr_max_requested_bandwidth_ul, args.apn2_ambr_max_requested_bandwidth_dl)
    ue_ambr = ambr2json( args.ue_ambr_max_requested_bandwidth_ul, args.ue_ambr_max_requested_bandwidth_dl)
    
    imsi_first_int = int(args.imsi_first)
    imsi_length = len(args.imsi_first)

    msisdn_first_int = int(args.msisdn_first)

    for i in range(args.no_of_users):
        imsi = imsi_first_int + i
        imsi_str = str(imsi)
        imsi_str = imsi_str.zfill(imsi_length)

        msisdn = msisdn_first_int + i

        # Do not throw away now
        #cmd_file = tempfile.NamedTemporaryFile(delete=False)
        #with open(cmd_file.name, "w") as f:
        #    f.write("cassandra_ip={0}".format(args.cassandra_cluster))
        #    f.write("imsi={0}".format(imsi_str))
        #    f.write("msisdn={0}".format(str(msisdn)))
        #    f.write("access_restriction=41")
        #    f.write("key={0}".format(args.key))
        #    f.write("mmehost={0}".format(args.mme_identity))
        #    f.write("mmeidentity_idmmeidentity={0}".format(3))
        #    f.write("mmerealm={0}".format(args.realm))
        #    f.write("rand={0}".format('2683b376d1056746de3b254012908e0e'))
        #    f.write("sqn={0}".format(96))
        #    f.write("""cqlsh $cassandra_ip -e "INSERT INTO vhss.users_imsi (imsi, msisdn, access_restriction, key, mmehost, mmeidentity_idmmeidentity, mmerealm, rand, sqn, subscription_data) VALUES ('$imsi', $msisdn, $access_restriction, '$key', '$mmehost', $mmeidentity_idmmeidentity, '$mmerealm', '$rand', $sqn, '{\"Subscription-Data\":{\"Access-Restriction-Data\":41,\"Subscriber-Status\":0,\"Network-Access-Mode\":2,\"Regional-Subscription-Zone-Code\":[\"0x0123\",\"0x4567\",\"0x89AB\",\"0xCDEF\",\"0x1234\",\"0x5678\",\"0x9ABC\",\"0xDEF0\",\"0x2345\",\"0x6789\"],\"MSISDN\":\"0x$msisdn\",\"AMBR\":{\"Max-Requested-Bandwidth-UL\":50000000,\"Max-Requested-Bandwidth-DL\":100000000},\"APN-Configuration-Profile\":{\"Context-Identifier\":0,\"All-APN-Configurations-Included-Indicator\":0,\"APN-Configuration\":{\"Context-Identifier\":0,\"PDN-Type\":0,\"Served-Party-IP-Address\":[\"10.0.0.1\",\"10.0.0.2\"],\"Service-Selection\":\"apn1\",\"EPS-Subscribed-QoS-Profile\":{\"QoS-Class-Identifier\":9,\"Allocation-Retention-Priority\":{\"Priority-Level\":15,\"Pre-emption-Capability\":0,\"Pre-emption-Vulnerability\":0}},\"AMBR\":{\"Max-Requested-Bandwidth-UL\":50000000,\"Max-Requested-Bandwidth-DL\":100000000},\"PDN-GW-Allocation-Type\":0,\"MIP6-Agent-Info\":{\"MIP-Home-Agent-Address\":[\"172.26.17.183\"]}}},\"Subscribed-Periodic-RAU-TAU-Timer\":0}}');" """)
        #    f.flush()

        served_party_ip_address = None

        if args.static_ue_ipv4_allocation is not None:
            try:
                ip4_addr = str(ipaddress.IPv4Address(args.static_ue_ipv4_allocation) + i)
                served_party_ip_address = '"Served-Party-IP-Address":["'+ip4_addr+'"'
            except:
                logging.exception("Bad parameter -s/--static-ue-ipv4-allocation")
                return os.EX_DATAERR

        if args.static_ue_ipv6_allocation is not None:
            try:
                ip6_addr = str(ipaddress.IPv6Address(args.static_ue_ipv6_allocation) + i)
                if served_party_ip_address is not None:
                    served_party_ip_address = served_party_ip_address+',"'+ip6_addr+'"]'
                else:
                    served_party_ip_address = '"Served-Party-IP-Address":["'+ip6_addr+'"]'
            except:
                logging.exception("Bad parameter -S/--static-ue-ipv6-allocation")
                return os.EX_DATAERR

        if served_party_ip_address is not None:

            session.execute(
                """
                INSERT INTO vhss.users_imsi (imsi, msisdn, access_restriction, key, opc, mmehost, mmeidentity_idmmeidentity, mmerealm, rand, sqn, subscription_data)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) 
                IF NOT EXISTS
                """,
                (imsi_str, msisdn, 41, args.key, args.opc, args.mme_identity, 3, args.realm, '2683b376d1056746de3b254012908e0e', args.sqn,
                '{"Subscription-Data":{"Access-Restriction-Data":41,"Subscriber-Status":0,"Network-Access-Mode":2,"Regional-Subscription-Zone-Code":["0x0123","0x4567","0x89AB","0xCDEF","0x1234","0x5678","0x9ABC","0xDEF0","0x2345","0x6789"],"MSISDN":"0x'+str(msisdn)+'",'+ue_ambr+',"APN-Configuration-Profile":{"Context-Identifier":0,"All-APN-Configurations-Included-Indicator":0,"APN-Configuration":{"Context-Identifier":0,"PDN-Type":0,'+served_party_ip_address+',"Service-Selection":"'+args.apn+'","EPS-Subscribed-QoS-Profile":{"QoS-Class-Identifier":9,"Allocation-Retention-Priority":{"Priority-Level":15,"Pre-emption-Capability":0,"Pre-emption-Vulnerability":0}},'+apn1_ambr+',"PDN-GW-Allocation-Type":0,"MIP6-Agent-Info":{"MIP-Home-Agent-Address":["172.26.17.183"]}},"APN-Configuration":{"Context-Identifier":1,"PDN-Type":0,'+served_party_ip_address+',"Service-Selection":"'+args.apn2+'","EPS-Subscribed-QoS-Profile":{"QoS-Class-Identifier":9,"Allocation-Retention-Priority":{"Priority-Level":13,"Pre-emption-Capability":1,"Pre-emption-Vulnerability":0}},'+apn2_ambr+',"PDN-GW-Allocation-Type":0,"MIP6-Agent-Info":{"MIP-Home-Agent-Address":["172.26.17.183"]}}},"Subscribed-Periodic-RAU-TAU-Timer":0}}'
                )
            )
        else:
            session.execute(
                """
                INSERT INTO vhss.users_imsi (imsi, msisdn, access_restriction, key, opc, mmehost, mmeidentity_idmmeidentity, mmerealm, rand, sqn, subscription_data)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                IF NOT EXISTS
                """,
                (imsi_str, msisdn, 41, args.key, args.opc, args.mme_identity, 3, args.realm, '2683b376d1056746de3b254012908e0e', args.sqn,
                '{"Subscription-Data":{"Access-Restriction-Data":41,"Subscriber-Status":0,"Network-Access-Mode":2,"Regional-Subscription-Zone-Code":["0x0123","0x4567","0x89AB","0xCDEF","0x1234","0x5678","0x9ABC","0xDEF0","0x2345","0x6789"],"MSISDN":"0x'+str(msisdn)+'",'+ue_ambr+',"APN-Configuration-Profile":{"Context-Identifier":0,"All-APN-Configurations-Included-Indicator":0,"APN-Configuration":{"Context-Identifier":0,"PDN-Type":0,"Service-Selection":"'+args.apn+'","EPS-Subscribed-QoS-Profile":{"QoS-Class-Identifier":9,"Allocation-Retention-Priority":{"Priority-Level":15,"Pre-emption-Capability":0,"Pre-emption-Vulnerability":0}},'+apn1_ambr+',"PDN-GW-Allocation-Type":0,"MIP6-Agent-Info":{"MIP-Home-Agent-Address":["172.26.17.183"]}},"APN-Configuration":{"Context-Identifier":1,"PDN-Type":0,"Service-Selection":"'+args.apn2+'","EPS-Subscribed-QoS-Profile":{"QoS-Class-Identifier":9,"Allocation-Retention-Priority":{"Priority-Level":13,"Pre-emption-Capability":1,"Pre-emption-Vulnerability":0}},'+apn2_ambr+',"PDN-GW-Allocation-Type":0,"MIP6-Agent-Info":{"MIP-Home-Agent-Address":["172.26.17.183"]}}},"Subscribed-Periodic-RAU-TAU-Timer":0}}'
                )
            )


        session.execute(
            """
            INSERT INTO vhss.msisdn_imsi (msisdn, imsi) VALUES (%s, %s) IF NOT EXISTS
            """,
            (msisdn, imsi_str))

    if str2bool(args.verbose):
        # TODO pretty print
        rows = session.execute('SELECT imsi, msisdn, access_restriction, key, mmehost, mmeidentity_idmmeidentity, mmerealm, rand, sqn, subscription_data FROM vhss.users_imsi')
        for row in rows:
            logging.info("%s %s %s %s %s %s %s %s %s %s" % (row[0], row[1], row[2], row[3], row[4], row[5], row[6], row[7], row[8], row[9]))

        rows = session.execute('SELECT msisdn, imsi FROM vhss.msisdn_imsi')
        for row in rows:
            logging.info("%s %s" % (row[0], row[1]))


if __name__ == "__main__":
    main()
